import cv2
import numpy as np
import matplotlib.pyplot as plt

img = cv2.imread("test.jpg")
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 转换为RGB用于matplotlib显示

hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

lower_red1 = np.array([0, 120, 70])    # 低饱和度红色区间
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 120, 70])  # 高饱和度红色区间
upper_red2 = np.array([180, 255, 255])

mask1 = cv2.inRange(hsv_img, lower_red1, upper_red1)
mask2 = cv2.inRange(hsv_img, lower_red2, upper_red2)
mask = cv2.bitwise_or(mask1, mask2)  # 合并两个红色区间的掩码

kernel = np.ones((5, 5), np.uint8)
mask = cv2.erode(mask, kernel, iterations=1)  # 腐蚀
mask = cv2.dilate(mask, kernel, iterations=1)  # 膨胀

result = cv2.bitwise_and(img, img, mask=mask)
result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(12, 6))
plt.subplot(1, 3, 1), plt.imshow(img_rgb), plt.title("原图"), plt.axis("off")
plt.subplot(1, 3, 2), plt.imshow(mask, cmap="gray"), plt.title("掩码（红色区域）"), plt.axis("off")
plt.subplot(1, 3, 3), plt.imshow(result_rgb), plt.title("颜色分割结果"), plt.axis("off")
plt.show()
